"""
K-means 算法演示
"""
import pandas as pd
import matplotlib.pyplot as plt

from k_means import KMeans

data = pd.read_csv('iris.csv')
iris_types = ['SETOSA', 'VERSICOLOR', 'VIRGINICA']

x_axis = 'petal_length'
y_axis = 'petal_width'

# 有无标签对比
plt.figure(figsize=(12, 5))
plt.subplot(121)
for iris_type in iris_types:
    plt.scatter(
        data[x_axis][data['class'] == iris_type],
        data[y_axis][data['class'] == iris_type],
        label=iris_type
    )
plt.legend()
plt.title('Label known')

plt.subplot(122)
plt.scatter(data[x_axis], data[y_axis])
plt.title('Label unknown')
plt.show()

# 进行 K-means 训练
num_examples = data.shape[0]
num_clusters = 3
max_iterations = 50
x_train = data[[x_axis, y_axis]].values.reshape(num_examples, 2)
k_means = KMeans(x_train, num_clusters)
centroids, closest_centroids_ids = k_means.train(max_iterations)

# K-means 训练效果展示
plt.figure(figsize=(12, 5))
plt.subplot(121)
for iris_type in iris_types:
    plt.scatter(
        data[x_axis][data['class'] == iris_type],
        data[y_axis][data['class'] == iris_type],
        label=iris_type
    )
plt.legend()
plt.title('Authentic')

plt.subplot(122)
for k, centroid in enumerate(centroids):
    plt.scatter(
        x_train[(closest_centroids_ids == k).flatten(), 0],
        x_train[(closest_centroids_ids == k).flatten(), 1],
        label=k
    )
    plt.plot([centroid[0]], [centroid[1]], 'rx')

plt.legend()
plt.title('K - means')
plt.show()
